from langchain_community.utilities import SQLDatabase
from mcp.server import FastMCP

from llm_utils import zhipuai_client

fastmcp_server = FastMCP(name='fw-mcp', instructions='这是测试MCP服务', port=6666)
db = SQLDatabase.from_uri('sqlite:///../chinook.db')


@fastmcp_server.tool(name='zhipu_search', description='搜索互联网中的内容')
def zhipu_search(query: str) -> str:
    """搜索互联网上的内容
    Args:
        query (str):要执行的查询信息

    Returns:
        str: 查询结果或错误信息
    """
    try:
        resp = zhipuai_client.web_search.web_search(
            query=query,
            search_engine='search-std'
        )
        print(f'智谱搜索结果：{resp}')
        if resp.search_result:
            return '\n'.join([item.title for item in resp.search_result])
    except Exception as e:
        print(f'智谱搜索异常：{e}')
        return '智谱搜索异常!!!'


@fastmcp_server.tool(name='list_tables_tool', description='返回数据库中所有的表，以逗号分隔表名')
def list_tables_tool() -> str:
    """返回数据库中所有的表,以逗号隔表名"""
    return ', '.join(db.get_usable_table_names())


@fastmcp_server.tool(name='db_query_tool', description='执行SQL查询并返回结果')
def db_query_tool(query: str) -> str:
    """执行SQL查询并返回结果
    Args:
        query (str):要执行的查询信息

    Returns:
        str: 查询结果或错误信息
    """
    result = db.run_no_throw(query) # 执行查询不抛出异常
    if not result:
        print(f'查询结果为空：{result}')
        return f"查询错误，请修改查询语句后重试，{query}"
    return result
